#!/usr/bin/env python3
"""
h264_to_fsv.py

将原始 H.264 Annex-B 流转换为 FreeSWITCH FSV 视频文件。
为解决播放速度问题，加入虚拟音频帧以驱动定时器。
"""

import struct
import sys
import time
import re
import argparse
import math

# FSV 常量
VID_BIT = 1 << 31
VERSION = 4202
SSRC_VIDEO = 0x12345678
SSRC_AUDIO = 0x87654321 # 虚拟音频 SSRC
PAYLOAD_TYPE_H264 = 96
PAYLOAD_TYPE_PCMU = 0   # 使用 PCMU (G.711 U-law) 作为虚拟音频
MTU = 1200
DEFAULT_FPS = 25
RTP_CLOCK_RATE_VIDEO = 90000
RTP_CLOCK_RATE_AUDIO = 8000 # PCMU 标准采样率
AUDIO_PTIME_MS = 20        # 音频包时长 (ms)
# 计算音频包的采样点数: 8000Hz * 20ms / 1000ms = 160 samples
AUDIO_SAMPLES_PER_PACKET = (RTP_CLOCK_RATE_AUDIO * AUDIO_PTIME_MS) // 1000 

def write_fsv_header(f, fps):
    """写入 FSV 文件头"""
    fmt = "<i32s128sII4xqI4x"
    version = VERSION
    # 视频编解码器
    video_codec_name = b"H264"
    video_fmtp = b""
    # 音频参数 - 对于 FSV 文件头是必需的，即使我们只处理视频
    # 我们将它们设置为常见的 PCMU 参数
    audio_rate = RTP_CLOCK_RATE_AUDIO 
    audio_ptime = AUDIO_PTIME_MS   
    created = int(time.time() * 1_000_000) # 微秒
    channels = 1 # 音频通道数
    
    frame_inc_video = RTP_CLOCK_RATE_VIDEO // fps
    print(f"Info: Target FPS={fps}")
    print(f"      Video TS increment per frame={frame_inc_video}")
    print(f"      Audio TS increment per packet={AUDIO_SAMPLES_PER_PACKET} samples ({AUDIO_PTIME_MS}ms)")

    packed = struct.pack(
        fmt, version, video_codec_name, video_fmtp,
        audio_rate, audio_ptime, created, channels
    )
    f.write(packed)
    # 返回视频帧时间戳增量
    return frame_inc_video

def build_rtp_header(seq, ts, ssrc, marker=0, payload_type=PAYLOAD_TYPE_H264):
    """构建标准 RTP 头部 (12 字节)，使用网络字节序 (!)。"""
    v_p_x_cc = 0x80  # Version=2, Padding=0, Extension=0, CC=0
    m_pt = ((marker & 0x1) << 7) | (payload_type & 0x7f)
    return struct.pack("!BBHII", v_p_x_cc, m_pt, seq & 0xFFFF, ts & 0xFFFFFFFF, ssrc & 0xFFFFFFFF)

def packetize_nalu(nalu, seq, ts, marker):
    """将单个 NALU 打包成一个或多个 RTP 包。"""
    packets = []
    if not nalu:
        return packets, seq

    # 1. 单包 RTP
    if len(nalu) + 12 <= MTU:
        rtp_packet = build_rtp_header(seq, ts, SSRC_VIDEO, marker, PAYLOAD_TYPE_H264)
        rtp_packet += nalu
        packets.append(rtp_packet)
        return packets, seq + 1

    # 2. FU-A 分片
    nalu_header = nalu[0]
    nalu_type = nalu_header & 0x1F
    fu_indicator = (nalu_header & 0xE0) | 28 # Type=28 for FU-A

    FU_A_HEADER_SIZE = 2
    max_fu_payload_size = MTU - 12 - FU_A_HEADER_SIZE

    if max_fu_payload_size <= 0:
        print(f"Warning: MTU ({MTU}) too small for FU-A. Skipping NALU.")
        return packets, seq

    offset = 1
    first = True
    while offset < len(nalu):
        take = min(max_fu_payload_size, len(nalu) - offset)
        frag = nalu[offset: offset + take]

        if first:
            fu_header = 0x80 | nalu_type # Start bit
            first = False
        elif offset + take >= len(nalu):
            fu_header = 0x40 | nalu_type # End bit
        else:
            fu_header = nalu_type # Middle fragment

        # Marker bit: Only set on the last FU-A packet of the NALU
        m = 1 if (offset + take >= len(nalu)) else 0
        # 如果这是整个逻辑帧的最后一个 NALU 的最后一个 FU-A 分片，则使用传入的 marker
        if m == 1 and marker:
             final_m = 1
        else:
             final_m = 0

        rtp_packet = build_rtp_header(seq, ts, SSRC_VIDEO, final_m, PAYLOAD_TYPE_H264)
        rtp_packet += bytes([fu_indicator, fu_header]) + frag
        packets.append(rtp_packet)
        seq += 1
        offset += take
        
    return packets, seq

def write_audio_frame(fout, seq_audio, ts_audio):
    """写入一个静音的虚拟音频帧 (PCMU)"""
    # 生成静音数据 (0xFF 对应 PCMU 的静音)
    # 160 个采样点，每个采样点 1 字节 (PCMU)
    silent_audio_data = b'\xFF' * AUDIO_SAMPLES_PER_PACKET 
    
    # 构建 RTP 头 (音频 PT=0, PCMU)
    # 音频帧通常不设置 Marker 位
    rtp_header_audio = build_rtp_header(seq_audio, ts_audio, SSRC_AUDIO, 0, PAYLOAD_TYPE_PCMU)
    audio_rtp_packet = rtp_header_audio + silent_audio_data
    
    # 写入 FSV 音频帧头 (长度，不带 VID_BIT)
    fout.write(struct.pack("<I", len(audio_rtp_packet) & 0x7FFFFFFF))
    # 写入完整的音频 RTP 包
    fout.write(audio_rtp_packet)
    
    # 返回更新后的序列号和时间戳
    return seq_audio + 1, ts_audio + AUDIO_SAMPLES_PER_PACKET

def parse_annexb(data):
    """解析 Annex-B 格式的 H.264 数据，提取 NALU"""
    matches = [m for m in re.finditer(b'\x00\x00\x01|\x00\x00\x00\x01', data)]
    nalus = []
    for i, m in enumerate(matches):
        sc_start = m.start()
        sc_len = 4 if data[sc_start:sc_start+4] == b'\x00\x00\x00\x01' else 3
        nalu_start = sc_start + sc_len
        if i + 1 < len(matches):
            nalu_end = matches[i+1].start()
        else:
            nalu_end = len(data)
        nalu = data[nalu_start:nalu_end]
        if nalu:
            nalus.append(nalu)
    return nalus

def split_frames(nalus):
    """将 NALU 列表按帧分组。"""
    frames = []
    current_frame = []
    for nalu in nalus:
        ntype = nalu[0] & 0x1F
        if ntype in (1, 5): # SLICE NALU 标识新帧
            if current_frame:
                frames.append(current_frame)
            current_frame = []
        current_frame.append(nalu)
    if current_frame:
        frames.append(current_frame)
    return frames

def h264_to_fsv(h264_file, fsv_file, fps):
    """主转换函数"""
    try:
        with open(h264_file, "rb") as fin:
            data = fin.read()
    except FileNotFoundError:
        print(f"Error: Input file '{h264_file}' not found.")
        return False
    except Exception as e:
        print(f"Error reading input file: {e}")
        return False

    nalus = parse_annexb(data)
    if not nalus:
        print("Error: No NALUs found in the input file.")
        return False

    print(f"Found {len(nalus)} NALUs in the input file.")

    frames = split_frames(nalus)
    if not frames:
        print("Error: Could not split NALUs into frames.")
        return False

    print(f"Grouped into {len(frames)} logical frames.")

    seq_video = 1
    seq_audio = 1
    ts_video = 0  # 视频帧时间戳
    ts_audio = 0  # 音频帧时间戳
    sps = None
    pps = None

    frame_inc_video = RTP_CLOCK_RATE_VIDEO // fps

    try:
        with open(fsv_file, "wb") as fout:
            write_fsv_header(fout, fps)

            audio_tick = 20
            video_tick = 1000 // fps
            current_tick = 0
            
            audio_cnt = 0
            video_cnt = 0

            video_fns = iter(frames)
            while True: 
                if current_tick % audio_tick == 0:
                    seq_audio, ts_audio = write_audio_frame(fout, seq_audio, ts_audio)
                    audio_cnt = audio_cnt + 1
                    pass
                if current_tick % video_tick == 0:
                    try: 
                        frame_nalus = next(video_fns)
                    except StopIteration:
                        break

                    # 1. 提取并缓存 SPS/PPS
                    for n in frame_nalus:
                        ntype = n[0] & 0x1F
                        if ntype == 7: sps = n
                        elif ntype == 8: pps = n

                    # 2. 确定此逻辑帧需要发送的 NALU 列表
                    nalus_to_send_with_marker = []
                    has_idr = any((n[0] & 0x1F) == 5 for n in frame_nalus)

                    if has_idr:
                        if sps:
                            nalus_to_send_with_marker.append((sps, False))
                        if pps:
                            nalus_to_send_with_marker.append((pps, False))

                    for j, nalu in enumerate(frame_nalus):
                        is_last_nalu_in_frame = (j == len(frame_nalus) - 1)
                        nalus_to_send_with_marker.append((nalu, is_last_nalu_in_frame))

                    # 3. 打包并写入所有 NALU (使用当前帧的 ts_video)
                    for k, (nalu, is_last_in_logical_frame) in enumerate(nalus_to_send_with_marker):
                        pkts, seq_video = packetize_nalu(nalu, seq_video, ts_video, 1 if is_last_in_logical_frame else 0)
                        for pkt in pkts:
                            fout.write(struct.pack("<I", (len(pkt) | VID_BIT) & 0xFFFFFFFF))
                            fout.write(pkt)

                    # 5. 处理完一个完整的逻辑帧后，递增视频时间戳
                    ts_video += frame_inc_video                    
                    video_cnt = video_cnt + 1
                current_tick = current_tick + 1

        print(f"Successfully converted '{h264_file}' to '{fsv_file}' at {fps}fps.")
        print(f"Final video timestamp: {ts_video}, Final audio timestamp: {ts_audio} ({ts_audio / RTP_CLOCK_RATE_AUDIO:.2f}s)")
        print(f"video_cnt: {video_cnt}, audio_cnt: {audio_cnt}")
        return True

    except Exception as e:
        print(f"Error writing output file: {e}")
        import traceback
        traceback.print_exc()
        return False

def main():
    parser = argparse.ArgumentParser(
        description="Convert raw H.264 Annex-B file to FreeSWITCH FSV format.",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument("input", help="Input .h264 file (Annex-B format)")
    parser.add_argument("output", help="Output .fsv file")
    parser.add_argument("--fps", type=int, default=DEFAULT_FPS,
                        help=f"Frame rate for the output FSV file (default: {DEFAULT_FPS})")

    args = parser.parse_args()

    if args.fps <= 0:
        print("Error: FPS must be a positive integer.")
        sys.exit(1)

    success = h264_to_fsv(args.input, args.output, args.fps)
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()